// 基本思路是：
// 1. 检查url是否在黑名单中，在则直接返回403
// 2. 提取每次请求的特征：header的长度+请求的uri
// 3. 是否已经有一个检查器在工作，若已存在，则计数器加1，若不存在，则新建一个检查器，并起协程，一个周期后卸载检查器
// 4. 检查器每次加一后检查是否已到阈值，若已到阈值，则添加请求ip到黑名单并卸载检查器
// 5. 周期到达后，卸载检查器

package main

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"log"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

var (
	blockedIPs  sync.Map
	checkersMap sync.Map // string->*checker
	featuresMap sync.Map // string->struct{}

)

const (
	threshold = uint64(500)
	period    = 1 * time.Second
)

type checker struct {
	feature string
	ip      string
	counter uint64
}

func (c *checker) add() {
	atomic.AddUint64(&(c.counter), 1)
	if c.counter > threshold {
		// 添加黑名单
		blockIp(c.ip)
		// 清除map
		unloadMap(c.feature)
	}
}

func (c *checker) unload() {
	if c.counter > threshold {
		blockIp(c.ip)
	}
	unloadMap(c.feature)

}

func main() {
	r := gin.New()
	r.Use(filter())
	r.Any("/", func(c *gin.Context) {
		c.JSON(http.StatusOK, "ok")
	})

	r.Run(":8090")
}

func filter() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 检查是否已在黑名单中
		ip := c.Request.RemoteAddr
		if checkBlock(ip) {
			c.AbortWithStatus(http.StatusForbidden)
			return
		}
		// 进入检查
		f := getFeature(c.Request)
		go checkReq(f, ip)

		c.Next()
	}
}

func checkBlock(ip string) bool {
	if _, ok := blockedIPs.Load(ip); ok {
		return true
	}
	return false
}

// 提取req特征（header长度+url）
func getFeature(r *http.Request) string {
	var length int
	for k, v := range r.Header {
		length += len(k) + len(v)
	}
	return fmt.Sprintf("%d:%s", length, r.RequestURI)
}

func checkReq(f, ip string) {
	var c *checker
	if _, ok := featuresMap.LoadOrStore(f, struct{}{}); ok {
		if ci, ok := checkersMap.Load(f); ok {
			//可能存在featuresMap 与 checkersMap 不同步的情况
			c = ci.(*checker)
			c.add()
		} else {
			c = &checker{feature: f, counter: 1, ip: ip}
			go time.AfterFunc(period, c.unload)
			checkersMap.Store(f, c)
		}
	} else {
		c = &checker{feature: f, counter: 1, ip: ip}
		go time.AfterFunc(period, c.unload)
		checkersMap.Store(f, c)
	}
}

func blockIp(ip string) {
	log.Printf("%s 已被添加到黑名单\n", ip)
	blockedIPs.Store(ip, struct{}{})
}

// 清空map
func unloadMap(f string) {
	checkersMap.Delete(f)
	featuresMap.Delete(f)
}
